module lang::ofg::ast::Java2OFG

import IO;
import Set;
import List;
import lang::ofg::ast::FlowLanguage;
import lang::java::m3::TypeSymbol;
import lang::java::jdt::m3::AST;

Program createOFG(loc project) = createOFG(createAstsFromEclipseProject(project, true));

Program createOFG(set[Declaration] asts) {
	println("Getting decls");
	decls = getDeclarations(asts);
	println("Getting stms");
	stms = getStatements(asts);
	return program(decls, stms);
}

set[str] containerClasses =  {
	 "/java/util/Map"
	,"/java/util/HashMap"
	,"/java/util/Collection"
	,"/java/util/Set"
	,"/java/util/HashSet"
	,"/java/util/LinkedHashSet"
	,"/java/util/List"
	,"/java/util/ArrayList"
	,"/java/util/LinkedList"
	// Extra containers:
	,"/java/util/Deque"
	,"/java/util/NavigableMap"
	,"/java/util/NavigableSet"
	,"/java/util/Queue"
	,"/java/util/SortedMap"
	,"/java/util/SortedSet"
	,"/java/util/AbstractCollection"
	,"/java/util/AbstractList"
	,"/java/util/AbstractMap"
	,"/java/util/AbstractQueue"
	,"/java/util/AbstractSequentialList"
	,"/java/util/AbstractSet"
	,"/java/util/ArrayDeque"
	,"/java/util/Bitset"
	,"/java/util/Dictionary"
	,"/java/util/EnumMap"
	,"/java/util/EnumSet"
	,"/java/util/Hashtable"
	,"/java/util/IdentityHashMap"
	,"/java/util/LinkedHashMap"
	,"/java/util/LinkedHashSet"
	,"/java/util/LinkedList"
	,"/java/util/PriorityQueue"
	,"/java/util/Stack"
	,"/java/util/TreeMap"
	,"/java/util/TreeSet"
	,"/java/util/Vector"
	,"/java/util/WeakHashMap"
	,"/java/util/concurrent/BlockingDeque"
	,"/java/util/concurrent/ConcurrentMap"
	,"/java/util/concurrent/ConcurrentNavigableMap"
	,"/java/util/concurrent/TransferQueue"
	,"/java/util/concurrent/ArrayBlockingQueue"
	,"/java/util/concurrent/ConcurrentHashMap"
	,"/java/util/concurrent/ConcurrentLinkedDeque"
	,"/java/util/concurrent/ConcurrentLinkedQueue"
	,"/java/util/concurrent/ConcurrentSkipListMap"
	,"/java/util/concurrent/ConcurrentSkipListSet"
	,"/java/util/concurrent/CopyOnWriteArrayList"
	,"/java/util/concurrent/CopyOnWriteArraySet"
	,"/java/util/concurrent/DelayQueue"
	,"/java/util/concurrent/LinkedBlockingDeque"
	,"/java/util/concurrent/LinkedBlockingQueue"
	,"/java/util/concurrent/LinkedTransferQueue"
	,"/java/util/concurrent/PriorityBlockingQueue"
	,"/java/util/concurrent/SynchronousQueue"
};

map[str, int] insertArgs = (
	  "insert": 0
	, "insertAll": 0
	, "put": 1
	, "putAll": 0
	, "add": 0
	, "addAll": 0
	// Extra insert methods:
	, "addFirst": 0 //1
	, "offerFirst": 0 //1
	, "addLast": 0 //1
	, "offerLast": 0 //1
	, "offer": 0 //1
	, "push": 0 //1
	//, "set": 0 //2 <-- for this one unsure what insertArg is the correct one.
);

Expression correctInsertArg(Expression recv, str name, list[Expression] args) {
	return args[insertArgs[name]];
}


bool isContainerInsert(Expression recv, str name) {
	// get the fully qualified type path
	tp = (recv@typ).decl.path;
	// is it defined in the container classes?
	if (tp in containerClasses) {
		// if so, then return true if it is also an insert method.
		return name in insertArgs;
	}
	return false;
}

bool isContainerExtract(Expression recv, str name) {
	// get the fully qualified type path
	tp = (recv@typ).decl.path;
	// is the type a container class?
	if (tp in containerClasses) {
		switch (name) {
			// if so, return true if it is a container extract method.
			case "get": return true;	
			case "iterator": return true;	
			case "toArray": return true;
			case "entrySet": return true;	
			case "values": return true;
			// Extra container extract methods:
			case "pollFirst": return true;
			case "peekFirst": return true;
		    case "removeFirst": return true;
		    case "getFirst": return true;
		    case "removeLast": return true;
		    case "pollLast": return true;
		    case "getLast": return true;
		    case "peekLast": return true;
		    case "remove": return true;
		    case "poll": return true;
		    case "element": return true;
		    case "peek": return true;
		    case "pop": return true;
		    case "listIterator": return true;
		    case "subList": return true;
		    case "next": return true;
		    case "previous": return true;
		    case "ceilingEntry": return true;
		    case "descendingMap": return true;
		    case "firstEntry": return true;
		    case "floorEntry": return true;
		    case "headMap": return true;
		    case "higherEntry": return true;
		    case "lastEntry": return true;
		    case "lowerEntry": return true;
		    case "pollFirstEntry": return true;
		    case "pollLastEntry": return true;
		    case "subMap": return true;
		    case "tailMap": return true;
		    case "ceiling": return true;
		    case "descendingIterator": return true;
		    case "descendingSet": return true;
		    case "floor": return true;
		    case "headSet": return true;
		    case "higher": return true;
		    case "lower": return true;
		    case "subSet": return true;
		    case "tailSet": return true;
		    case "entrySet": return true;
		    case "first": return true;
		    case "last": return true;
		}
	}
	return false;
}

set[Declaration] fixCollections(set[Declaration] ast) {
	// perform a tree visit
	return visit (ast) {
		// find method calls
		case oe:methodCall(_, Expression receiver, methodName,	args):  {
			// is the methodcall a container insert? Ex: a.insert(b) where "a" is a container and "insert" is an insert method
			if (isContainerInsert(receiver, methodName)) {
				// correctInsertArg finds the correct parameter that is actually inserted into the container, which is "b".
				// insert a new assignment of the form a = b; where a is the collection and b is the correct insert arg.
				// In short, "a.insert(b)" is replaced with "a = b". The algorithm will then correctly see that data flows
				// from container b to a.
				insert assignment(receiver, "=", correctInsertArg(receiver, methodName, args))
					[@typ = receiver@typ] // the type of the assignment is the receiver type (type of the collection a)
					[@src = oe@src]; // the source of the assignment is the original method's source
			}
			// if not, is the methodcall a container extract? Ex: a = b.pop() or someOtherCall(b.pop()) where b is a container and "pop" is an extract method.
			else if(isContainerExtract(receiver, methodName)) {
				// if so, then replace "a = b.pop()" with "a = b", where b is a container.
				// or, as in the second example, replace "someOtherCall(b.pop())" with "someOtherCall(b)", where b is a container.
				// The algorithm will then correctly see the data flow from container b to the object a or to some method.
				insert receiver;
			}
		}
	};
}

set[str] primitiveTypes = {
	 "Byte" ,"java.lang.Byte"
	,"Character" ,"java.lang.Character"
	,"Short" ,"java.lang.Short"
	,"Integer" ,"java.lang.Integer"
	,"Long" ,"java.lang.Long"
	,"Float" ,"java.lang.Float"
	,"Double" ,"java.lang.Double"
	// Extra primitive types:
	,"AtomicInteger" ,"java.util.concurrent.atomic.AtomicInteger"
	,"AtomicLong" ,"java.util.concurrent.atomic.AtomicLong"
	,"BigDecimal" ,"java.math.BigDecimal"
	,"BigInteger" ,"java.math.BigInteger"
	,"Boolean" ,"java.lang.Boolean"
	,"String" ,"java.lang.String"
};
bool ignoreType(arrayType(t)) = ignoreType(t);
bool ignoreType(upperbound(t)) = ignoreType(t);
bool ignoreType(lowerbound(t)) = ignoreType(t);
bool ignoreType(parameterizedType(_)) = false; // A parameterized type is a generic, such as Collection<User>
bool ignoreType(Type::qualifiedType(t,_)) = ignoreType(t);
bool ignoreType(simpleType(t)) = ignoreType(t);// t in primitiveTypes;
bool ignoreType(unionType(tt)) = (false | it || ignoreType(t) | t <- tt);
default bool ignoreType(Type t) = true;

bool ignoreType(Expression::simpleName(n)) = n in primitiveTypes;
bool ignoreType(Expression::qualifiedName(q,n)) {
	if (simpleName(nn) := n && nn in primitiveTypes) {
		// could be a primitive type	
		// lets build the fullTypeName
		result = nn;
		parent = q;
		while (true) {
			if (qualifiedName(p,simpleName(c)) := parent) {
				result = c + "." + result;	
				parent = p;
			}	
			else if (simpleName(c) := parent) {	
				result = c + "." + result;
				break;
			}
			else {
				throw "unexpected parent c";	
			}
		}
		return result in primitiveTypes;
	}
	else {
		return false;	
	}
}
default bool ignoreType(Expression e) { throw "you forgot: <e>"; }


bool ignoreType(TypeSymbol::interface(_,_)) = false; // loc and list of types
bool ignoreType(TypeSymbol::\enum(_)) = false; // loc
bool ignoreType(TypeSymbol::\typeParameter(_,_)) = false; // loc and upperbound
bool ignoreType(TypeSymbol::\wildcard(_)) = false; // bound
bool ignoreType(TypeSymbol::\capture(_,_)) = false; //bound and type
bool ignoreType(TypeSymbol::intersection(tt)) = (false | it || ignoreType(t) | t <- tt); // list of types
bool ignoreType(TypeSymbol::union(tt)) = (false | it || ignoreType(t) | t <- tt); // list of types
bool ignoreType(TypeSymbol::\class(t,_)) = t == |java+class:///java/lang/String|; // loc and list of types
bool ignoreType(TypeSymbol::\object()) = false;
bool ignoreType(TypeSymbol::\array(_,_)) = true; // type and dimension
default bool ignoreType(TypeSymbol t) = true;

// Gets all declarations (attributes, methods and constructors) of types that are not ignored.
set[Decl] getDeclarations(set[Declaration] asts) 
	= { Decl::attribute(v@decl) | /field(t,frags) <- asts, !ignoreType(t), v <- frags}
	+ { Decl::method(m@decl, [p@decl | p:parameter(t,_,_) <- params, !ignoreType(t)]) | /m:Declaration::method(_,_, list[Declaration] params, _, _)  <- asts}
	+ { Decl::method(m@decl, [p@decl | p:parameter(t,_,_) <- params, !ignoreType(t)]) | /m:Declaration::method(_,_, list[Declaration] params, _)  <- asts}
	+ { Decl::constructor(c@decl, [p@decl | p:parameter(t,_,_) <- params, !ignoreType(t)]) | /c:Declaration::constructor(_, list[Declaration] params, _,_)  <- asts}      
	// add implicit constructor
	+ { Decl::constructor((c@decl)[scheme="java+constructor"] + "<name>()", []) | /c:class(name, _, _, b) <- asts, !(Declaration::constructor(_, _, _, _) <- b)}   
	;

loc lhsDecl(arrayAccess(e,_)) = e@decl;
loc lhsDecl(f:fieldAccess(_,_,_)) = f@decl;
loc lhsDecl(f:fieldAccess(_,_)) = f@decl;
loc lhsDecl(v:variable(_,_)) = v@decl;
loc lhsDecl(s:simpleName(_)) = s@decl;
loc lhsDecl(q:qualifiedName(_,_)) = q@decl;
default loc lhsDecl(Expression e) { throw "forgot: <e>"; }

set[Stm] getStatements(set[Declaration] asts) {
	// Find all methods.
	allMethods 
		= { m | /m:Declaration::method(_,_,_,_,_) <- asts}
		+ {Declaration::method(t, n, p, e, empty())[@decl=m@decl] | /m:Declaration::method(Type t,n,p,e) <- asts} 
		+ {Declaration::method(simpleType(simpleName(n)), n, p, e, b)[@decl=m@decl] | /m:Declaration::constructor(str n,p,e, b) <- asts} 
	;
	// Replace all collection insert and extract calls with ordinary statements.
	allMethods = fixCollections(allMethods);
	// now remove all nested classes to make all statements relative to a method
	allMethods = visit(allMethods) {
		case declarationExpression(Declaration::class(_)) => Expression::null()
		case declarationExpression(Declaration::class(_,_,_,_)) => Expression::null()
		case declarationExpression(Declaration::enum(_,_,_,_)) => Expression::null()
		case declarationStatement(Declaration::class(_)) => empty()
		case declarationStatement(Declaration::class(_,_,_,_)) => empty()
		case declarationStatement(Declaration::enum(_,_,_,_)) => empty()
	};
	
	set[Stm] result = {};
	
	for (m:Declaration::method(_, _, _, _, b) <- allMethods) {
		top-down-break visit(b) {
			// Add return statements.
			case \return(e) : 
				result += { *translate(m@decl, m@decl + "return", e)};
			// Add assignments and nested assignments.
			case e:Expression::assignment(l,_,r) : 
				if (!ignoreType(e@typ)) {
					result += { *translate(m@decl, lhsDecl(l), r)};
				} else {
					// there can be a nested assignment caused by the rewriting done earlier (containers)
					for (/e2:assignment(l2,_,r2) := r && !ignoreType(e2@typ)) {
						result += { *translate(m@decl, lhsDecl(l2), r2)};
					}
				}
			// Special assignments, like b = someMethod() or b = (castToType)SomeMethod() or b = new SomeObject()
			case v:Expression::variable(_,_,r) : 
				if (!ignoreType(v@typ)) {
					result += { *translate(m@decl, v@decl, r)};
				}
			// regular method calls with no target
			case m2:Expression::methodCall(_ ,_, _):
				result += { *translate(m@decl, emptyId, m2)};
			case m2:Expression::methodCall(_ ,_, _, _):
				result += { *translate(m@decl, emptyId, m2)};
		}
	}
	
	// Attempt to build in native array support. UNCOMPLETE!
	fieldVars = { f | /field(t,frags) <- asts, !ignoreType(t), f:Expression::variable(_,_,_) <- frags };
	newArrayTypes = { <f@decl, tp> | f <- fieldVars, v:Expression::newArray(tp,_) <- f };
	for (<frag, tp> <- newArrayTypes)
	{
		simpleType(sn) = tp;
		result += {Stm::assign(frag, emptyId, sn@decl)};
	}
	newArrayTypes = { <f@decl, tp> | f <- fieldVars, v:Expression::newArray(tp,_,_) <- f };
	for (<frag, tp> <- newArrayTypes)
	{
		simpleType(sn) = tp;
		result += {Stm::assign(frag, emptyId, sn@decl)};
	}
	
	return result;
}

// Translate returns a set of statements: assignments, new, call, newAssign, basically all data flows.
// TODO: handle a.b.c => B.c
set[Stm] translate(loc base, loc target, c:cast(_, e)) {
	// return the empty set when the type is ignored
	if (ignoreType(c@typ)) return {};
	// recursively call translate
	result = translate(base, target, e);
	// ??? returns result when s.target != target
	return { s.target == target ? s[cast=c@typ.decl] : s | s <- result};
}

set[Stm] translate(loc base, loc target, conditional(con, t, e)) 
	= translate(base, emptyId, con)
	+ translate(base, target, t)
	+ translate(base, target, e)
	;
	
// TODO: check what the second argument could mean (Expr)
set[Stm] translate(loc base, loc target, f:fieldAccess(_,_,_))
	= {Stm::assign(target, emptyId, f@decl)};
set[Stm] translate(loc base, loc target, f:fieldAccess(_,_))
	= {Stm::assign(target, emptyId, f@decl)};

set[Stm] translate(loc base, loc target, s:simpleName(_))
	= {Stm::assign(target, emptyId, s@decl)};

// nested assignment a = b = c;
set[Stm] translate(loc base, loc target, a:assignment(l,_,r)) 
	= translate(base, target, l)
	+ translate(base, target, r)
	;

set[Stm] translate(loc base, loc target, m:methodCall(s, n, a))
	= translate(base, target, methodCall(s, this(), n, a)[@decl=m@decl][@typ=m@typ][@src=m@src]);
set[Stm] translate(loc base, loc target, m:methodCall(_, r, n, a)) {
	set[Stm] stms = {};
	loc recv = emptyId;
	if (this() := r) {
		recv = base+"this";	
	}
	else {
		<newId, newStms> = unnestExpressions(base, r@src.offset, [r]);
		if (size(newId) > 0) {
			assert size(newId) == 1;
			recv = getOneFrom(newId);
		}
		stms += newStms;
	}
	<args, newStms> = unnestExpressions(base, m@src.offset, a);
	return newStms + { Stm::call(target, emptyId, recv, m@decl, args) };
}

private Expression newObject(Type t, list[Expression] args, Expression original) {
	assert original is newObject;
	return newObject(t, args)
		[@typ = original@typ]
		[@src = original@src]
		[@decl = original@decl];
}

set[Stm] translate(loc base, loc target, ob:newObject(_, Type t, a))
	= translate(base, target, newObject(t, a, ob));
set[Stm] translate(loc base, loc target, ob:newObject(_, Type t, a, _))
	= translate(base, target, newObject(t, a, ob));
set[Stm] translate(loc base, loc target, ob:newObject(Type t, a,_))
	= translate(base, target, newObject(t, a, ob));
set[Stm] translate(loc base, loc target, ob:newObject(Type t, a)) {
	assert target != emptyId;
	if (ignoreType(ob@typ))
		return {};
	
	<args, stms> = unnestExpressions(base, ob@src.offset, a);
	return stms + { Stm::newAssign(target, ob@typ.decl, ob@decl, args)};
}

bool simpleExpression(fieldAccess(_,_,_)) = true;
bool simpleExpression(fieldAccess(_,_)) = true;
bool simpleExpression(qualifiedName(_,e)) = simpleExpression(e);
bool simpleExpression(this()) = true;
bool simpleExpression(this(_)) = true;
bool simpleExpression(simpleName(_)) = true;
default bool simpleExpression(Expression e) = false;

Expression removeNesting(cast(_, e)) = removeNesting(e);
Expression removeNesting(arrayAccess(e, _)) = removeNesting(e);
Expression removeNesting(\bracket(e)) = removeNesting(e);
default Expression removeNesting(Expression e) = e;

// for arguments we have to unnestExpressions
//  .. = new A(new B());
// becomes
// __param<unique>_0 = new B();
// .. = new A(__param<unique>_0);
tuple[list[loc], set[Stm]] unnestExpressions(loc prefix, int uniqNum, list[Expression] exprs) {
	list[loc] ids = [];
	set[Stm] newStms = {};
	for (i <- [0..size(exprs)], Expression ce := exprs[i], !ignoreType(ce@typ)) {
		ce = removeNesting(ce);
		if (simpleExpression(ce)) {
			if (ce is this) {
				ids += [prefix + "this"];
			} 
			else {
				ids += [ce@decl];
			}
		}
		else {
			newId = prefix + "__param<uniqNum>_<i>";
			ids += [newId];
			newStms += translate(prefix, newId, ce);
		}
	}
	return <ids, newStms>;
}

default set[Stm] translate(loc base, loc target, Expression e) = { *translate(base, target, ch) | Expression ch <- e};

